import h5py
import numpy as np
import scipy.io as sio
import torch
from collections import Counter


def load_data(private_dataset, opt):
    '''
    加载并划分数据集
    '''
    domains_list = private_dataset.DOMAINS_LIST
    domains_len = len(domains_list)

    if opt.rand_dataset:
        max_num = 10
        is_ok = False

        while not is_ok:
            if opt.dataset == 'fl_officecaltech':
                selected_domain_list = np.random.choice(domains_list, size=opt.num_users - domains_len, replace=True, p=None)
                selected_domain_list = list(selected_domain_list) + domains_list
            elif opt.dataset == 'fl_digits':
                selected_domain_list = np.random.choice(domains_list, size=opt.num_users, replace=True, p=None).tolist()

            result = dict(Counter(selected_domain_list))

            for k in result:
                if result[k] > max_num:
                    is_ok = False
                    break
            else:
                is_ok = True

    else:
        selected_domain_dict = {'mnist': 3, 'usps': 2, 'svhn': 1, 'syn': 4}  # 10

        selected_domain_list = []
        for k in selected_domain_dict:
            domain_num = selected_domain_dict[k]
            for i in range(domain_num):
                selected_domain_list.append(k)

        selected_domain_list = np.random.permutation(selected_domain_list).tolist()
        result = Counter(selected_domain_list)
    print(f'客户端数据集选择：{result},{selected_domain_list}')

    pri_train_loaders, test_loaders = private_dataset.get_data_loaders(selected_domain_list)

    return pri_train_loaders, test_loaders